/**
 * @file encrypt_cls2_o.cpp
 * @brief 用于加密和解密ENNX模型文件的实现
 * @details 该文件包含了用于处理ENNX模型文件的加密和解密功能,使用AES-CBC模式进行加密。
 *          主要功能包括:
 *          - ENNX文件的解密
 *          - 文件数据的读取和处理
 *          - 模型转换和加密
 * @author 
 * @date 2024-01
 */


#include "encrypt_cls2_o.h"                                  // 包含 YOLO v6 类定义


//#include "dll.h"                                      // DLL 库（具体功能未在代码中显示）
using namespace CryptoPP;                             // 使用 Crypto++ 命名空间

/**
 * @brief Decrypts ENNX data using AES-CBC mode
 * @param v_encrypt Encrypted data vector
 * @param key Encryption key
 * @param iv Initialization vector
 * @return Decrypted data vector
 */
std::vector<char> decryptENNX(
	const std::vector<char> & v_encrypt,
	const std::string& key,
	const std::string& iv) {
	std::vector<char> v_decrypt;                            // Store decrypted result
	try {
		CBC_Mode<AES>::Decryption decryption(reinterpret_cast<const byte *>(key.data()), key.size(), (const byte*)iv.data());

		std::vector<char> _decrypt = v_encrypt;             // Copy encrypted data
		decryption.ProcessData(
			reinterpret_cast<CryptoPP::byte*>(_decrypt.data()),
			reinterpret_cast<const CryptoPP::byte*>(v_encrypt.data()),
			v_encrypt.size());

		byte paddingByte = _decrypt.back();                 // Get padding byte
		size_t paddingSize = static_cast<size_t>(paddingByte);
		_decrypt.resize(v_encrypt.size() - paddingSize);    // Remove padding

		v_decrypt = _decrypt;
	} catch (const CryptoPP::Exception& e) {
		std::cerr << e.what() << std::endl;
	}
	return v_decrypt;
}


/**
 * @brief Reads file data and returns a character vector
 * @param file_path Path to the file
 * @return Vector containing the file contents
 */
// Read file contents into vector <char>
std::vector<char> get_file_data(const std::string& file_path) {
	std::ifstream file(file_path, std::ios_base::binary);
	if (!file.is_open()) {
		throw std::runtime_error("Failed to open file: " + file_path);
	}
	file.seekg(0, std::ios_base::end);
	const size_t file_len = file.tellg();
	file.seekg(0, std::ios_base::beg);
	std::vector<char> data(file_len, 0);
	if (file.read(data.data(), file_len)) {
		return data;
	} else {
		throw std::runtime_error("Failed to read file: " + file_path);
	}
}
/**
 * @brief Writes vector data to file
 * @param data Data to write
 * @param file_path Path to output file
 */
void write_file_data(std::vector<char> data, std::string file_path) {
	std::ofstream file(file_path, std::ios_base::binary);   // Open file in binary mode
	file.write(data.data(), data.size());                   // Write vector data to file
	file.close();                                           // Close the file
}
/**
 * Performs a simple encryption operation on the data.
 * @param data Data to be encrypted.
 */
void encrypt_data(std::vector<char>& data) {
	char key = 0x1; // Simple key, should be more complex in practice
	for (size_t i = 0; i < data.size(); ++i) {
		data[i] ^= key; // XOR operation for encryption
	}
}

/**
 * Performs a simple decryption operation on the data.
 * @param data Data to be decrypted.
 */
void decrypt_data(std::vector<char>& data) {
	char key = 0x1; // Same key as encryption
	for (size_t i = 0; i < data.size(); ++i) {
		data[i] ^= key; // XOR operation for decryption
	}
}
/**
 * @brief Merges XML and BIN files into a single encrypted output file
 * @param xml_path Path to the XML file
 * @param bin_path Path to the BIN file 
 * @param info Structure containing model information
 * @param output_path Path to the output merged file
 * @throws std::runtime_error If file operations fail
 */
void merge_files(
    const std::string& xml_path,
    const std::string& bin_path,
    const st_encrypt_info_ovino& info,
    const std::string& output_path) {
    std::vector<char> xml_data = get_file_data(xml_path);
    std::vector<char> bin_data = get_file_data(bin_path);

    // 加密XML和BIN数据
    encrypt_data(xml_data);
    encrypt_data(bin_data);

    // 创建头部信息
    st_encrypt_header header = {
        sizeof(st_encrypt_info_ovino), // 结构体大小
        xml_data.size(),               // XML文件大小
        bin_data.size()                // BIN文件大小
    };

    // 将结构体和文件数据合并到一个 vector 中
    std::vector<char> merged_data;
    merged_data.reserve(sizeof(st_encrypt_header) + sizeof(st_encrypt_info_ovino) + xml_data.size() + bin_data.size());
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&header), reinterpret_cast<const char*>(&header) + sizeof(st_encrypt_header));
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&info), reinterpret_cast<const char*>(&info) + sizeof(st_encrypt_info_ovino));
    merged_data.insert(merged_data.end(), xml_data.begin(), xml_data.end());
    merged_data.insert(merged_data.end(), bin_data.begin(), bin_data.end());

    // 写入文件
    std::ofstream output_file(output_path, std::ios_base::binary);
    if (!output_file.is_open()) {
        throw std::runtime_error("Failed to open output file: " + output_path);
    }
    output_file.write(merged_data.data(), merged_data.size());
    if (!output_file) {
        throw std::runtime_error("Failed to write to output file: " + output_path);
    }
}

/**
 * @brief Loads data from YAML config file into st_encrypt_info struct
 * @param input Conversion input structure containing config file path and other info
 * @return Populated st_encrypt_info struct
 * @throws std::runtime_error If config file has issues or reading fails
 */
st_encrypt_info_ovino loadConfigToInfo(
    const st_trans_model_input& input) 
	{
    YAML::Node root = YAML::LoadFile(input.cfg_path);
    if (!root["CLS"].IsDefined()) {
        throw std::runtime_error("Node not found: CLS");
    }
    YAML::Node config = root["CLS"];

    st_encrypt_info_ovino info;

    std::strcpy(info.gpu_model, input.gpu_model.c_str());
    info.batch_size = config["BATCH_SIZE"].as<int>();
    info.input_channel = config["INPUT_CHANNEL"].as<int>();

    info.image_width = config["IMAGE_WIDTH"].as<int>();
	info.image_height = config["IMAGE_HEIGHT"].as<int>();

    std::vector<float> img_mean = config["img_mean"].as<std::vector<float>>();
    std::vector<float> img_std = config["img_std"].as<std::vector<float>>();
    if (img_mean.size() != 3 || img_std.size() != 3) {
        throw std::runtime_error("Invalid img_mean or img_std configuration");
    }
    for (int i = 0; i < 3; ++i) {
        info.img_mean[i] = img_mean[i];
        info.img_std[i] = img_std[i];
    }

    std::string image_order = config["image_order"].as<std::string>();
    std::strcpy(info.image_order, image_order.c_str());

    info.dynamic = config["Dynamic"].as<bool>();

    return info;
}

/**
 * @brief Load and decrypt model file (supports .onnx and .ennx files)
 * @param input Conversion input structure containing model path and other info
 * @param key Decryption key
 * @param iv Initialization vector for decryption
 * @return Byte vector containing model file data
 * @throws std::runtime_error If file format is unsupported or reading fails
 */
std::vector<char> loadAndDecrypt_ONNX_ENNX_Model(
	const st_trans_model_input& input,
	const std::string& key,
	const std::string& iv)
{
	namespace fs = boost::filesystem;

	// Check model file extension
	fs::path p(input.model_path.c_str());
	if (p.extension().string() == ".onnx") {
		// For `.onnx` model file, read data directly
		return get_file_data(p.string());
	} else if (p.extension().string() == ".ennx") {
		// For encrypted `.ennx` model file
		std::vector<char> v_e_model = get_file_data(p.string()); // Load encrypted model file
		return decryptENNX(v_e_model, key, iv);                  // Decrypt data using key and IV
	} else {
		// Unsupported file format
		throw std::runtime_error("Unsupported model file format: " + p.extension().string());
	}
}

/**
 * @brief 删除临时模型文件
 * @param directory 要搜索的目录路径,默认为当前目录"."
 * @details 遍历指定目录,删除所有以"temp_model."开头的临时文件
 * @note 删除失败时会输出错误信息但不会中断执行
 */
void remove_temp_model_files(const std::string& directory = ".") {
	for (const auto& entry : std::filesystem::directory_iterator(directory)) {
		if (entry.is_regular_file()) {
			std::string filename = entry.path().filename().string();
			if (filename.rfind("temp_model.", 0) == 0) { // 检查文件名是否以 "temp_model." 开头
				try {
					std::filesystem::remove(entry.path());
					// std::cout << "Deleted: " << entry.path() << std::endl;
				} catch (const std::filesystem::filesystem_error& e) {
					std::cerr << "Error deleting " << entry.path() << ": " << e.what() << std::endl;
				}
			}
		}
	}
}
/**
 * @brief Convert ONNX model data to IR format XML and BIN file data buffers
 * @param core OpenVINO core object
 * @param model_data ONNX model data as string
 * @return A pair of buffers containing XML and BIN file data
 */
std::pair<std::vector<char>,std::vector<char>>
coreModel_to_XB_buffer(
	ov::Core& core,
	const std::string& model_data)
{
	std::shared_ptr<ov::Model> model = core.read_model(model_data, ov::Tensor{});

	const std::string temp_xml_path = "temp_model.xml";
	const std::string temp_bin_path = "temp_model.bin";

	ov::serialize(model, temp_xml_path, temp_bin_path);

	std::ifstream xml_file(temp_xml_path, std::ios::binary);
	std::ifstream bin_file(temp_bin_path, std::ios::binary);

	std::vector<char> xml_buffer((std::istreambuf_iterator<char>(xml_file)), std::istreambuf_iterator<char>());
	std::vector<char> bin_buffer((std::istreambuf_iterator<char>(bin_file)), std::istreambuf_iterator<char>());

	std::cout << "XML data size: " << xml_buffer.size() << " bytes" << std::endl;
	std::cout << "BIN data size: " << bin_buffer.size() << " bytes" << std::endl;

	std::remove(temp_xml_path.c_str());
	std::remove(temp_bin_path.c_str());

	return {xml_buffer, bin_buffer};
}

/**
 * @brief 将加密后的XML和BIN数据写入到hdats文件
 * @param xml_data XML数据缓冲区,将被加密
 * @param bin_data BIN数据缓冲区,将被加密
 * @param info 模型相关信息结构体
 * @param output_path 输出文件路径
 * @throw std::runtime_error 如果文件操作失败则抛出异常
 */
void encrypt_XML_BIN_to_hdats(
    std::vector<char>& xml_data,
    std::vector<char>& bin_data,
    const st_encrypt_info_ovino& info,
    const std::string& output_path)
{

    std::cout << std::endl;
    // 加密XML和BIN数据
    encrypt_data(xml_data);
    encrypt_data(bin_data);

    // 创建头部信息
    st_encrypt_header header = {
        sizeof(st_encrypt_info_ovino), // 结构体大小
        xml_data.size(),               // XML文件大小
        bin_data.size()                // BIN文件大小
    };

    // 将结构体和文件数据合并到一个 vector 中
    std::vector<char> merged_data;
    merged_data.reserve(sizeof(st_encrypt_header) + sizeof(st_encrypt_info_ovino) + xml_data.size() + bin_data.size());
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&header), reinterpret_cast<const char*>(&header) + sizeof(st_encrypt_header));
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&info), reinterpret_cast<const char*>(&info) + sizeof(st_encrypt_info_ovino));
    merged_data.insert(merged_data.end(), xml_data.begin(), xml_data.end());
    merged_data.insert(merged_data.end(), bin_data.begin(), bin_data.end());

    // 写入文件
    std::ofstream output_file(output_path, std::ios_base::binary);
    if (!output_file.is_open()) {
        throw std::runtime_error("Failed to open output file: " + output_path);
    }
    output_file.write(merged_data.data(), merged_data.size());
    if (!output_file) {
        throw std::runtime_error("Failed to write to output file: " + output_path);
    }
}

/**
 * @brief modeltrans_segv2 析构函数
 */
modeltrans_segv2::modeltrans_segv2() = default;

/**
 * @brief modeltrans_segv2 析构函数
 */
modeltrans_segv2::~modeltrans_segv2() = default;

/**
 * @brief encrpyt dlm model file
 * @param input Input structure containing config file path and other information
 * @return Returns true if conversion succeeds, false otherwise
 */
bool modeltrans_segv2::TransModel(const st_trans_model_input& input)
{
	// Model encryption parameters
	st_encrypt_info_ovino info;
	// ONNX model cache vector
	std::vector<char> v_model;

	try {
		info= loadConfigToInfo(input);
		// Use info in subsequent code
	} catch (const std::runtime_error& e) {
		std::cerr << "Error loading configuration: " << e.what() << std::endl;
		return false;
	}
	try {
		const std::string key = "irkdh_573%3?iq5h";   // Decryption key
		const std::string iv = "1234567890123456";    // Init vector
		v_model = loadAndDecrypt_ONNX_ENNX_Model(input, key, iv);
		// Use v_model in subsequent code
	}
	catch (const std::runtime_error& e) {
		std::cerr << "Error loading model: " << e.what() << std::endl;
		return false;
	}

	// Convert ONNX model to OpenVINO IR engine model
	// Initialize OpenVINO core object
	ov::Core core;
	// Convert vector<char> to string for read_model input
	const std::string model_data(v_model.begin(), v_model.end());

	// Convert model to buffer
	const auto [xml_, bin_] =
		coreModel_to_XB_buffer(core, model_data);

	// Get XML and BIN data buffers
	std::vector<char> xml_buffer = xml_;
	std::vector<char> bin_buffer = bin_;

	// Encrypt model and model info
	// Fill engine data info
	info.engine_size = xml_buffer.size() + bin_buffer.size();
	info.gpu_model[0] = *"Intel core";          // Init GPU model string

	// Get encrypted file save path
	const std::string encrypt_file = input.trans_path;
	// Call encryption function
	encrypt_XML_BIN_to_hdats(
		xml_buffer,
		bin_buffer,
		info,
		encrypt_file);

	remove_temp_model_files();
	// Return true on success
	return true;
}
